import importlib

import torch
from packaging import version
from transformers.integrations.tensor_parallel import ParallelInterface
from transformers.utils.import_utils import _is_package_available

#pip install torch>=2.5.0 torchvision>=0.15.0 torchaudio>=2.5.0 --index-url https://download.pytorch.org/whl/cu121

def is_torch_greater_or_equal(library_version: str, accept_dev: bool = False):
    """
    Accepts a library version and returns True if the current version of the library is greater than or equal to the
    given version. If `accept_dev` is True, it will also accept development versions (e.g. 2.7.0.dev20250320 matches
    2.7.0).
    """
    if not _is_package_available("torch"):
        print("_is_package_available")
        return False
    if accept_dev:
        print("accept_dev")
        return version.parse(version.parse(importlib.metadata.version("torch")).base_version) >= version.parse(library_version)
    else:
        print("else")
        return version.parse(importlib.metadata.version("torch")) >= version.parse(library_version)
_torch_distributed_available = torch.distributed.is_available()
if is_torch_greater_or_equal("2.5") and _torch_distributed_available:
    ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface()
else:
    ALL_PARALLEL_STYLES = None
    print(is_torch_greater_or_equal("2.5"))
    print( _torch_distributed_available)
print(ALL_PARALLEL_STYLES)
